from random import randint
import json
import bpy
import subprocess
import numpy as np
import os
import math

output_folder_prefix = './tablescene'
scene_type = 'DiningRoom'

def random_hex(length=16):
    return "".join([hex(randint(0, 15))[2:] for _ in range(length)])

def to_hex(x: int):
    return hex(x)[2:]

def get_scene_path(id):
    return f'./tablescene/DiningRoom_{id}/coarse/scene.blend'
def get_scene_folder(id):
    return f'./tablescene/DiningRoom_{id}'

def get_image_path(id):
    return f'./tablescene/DiningRoom_{id}/frames/Image/camera_0/Image_0_0_0001_0.png'

def get_depth_path(id):
    return f'./tablescene/DiningRoom_{id}/frames/Depth/camera_0/Depth_0_0_0001_0.npy'

def get_objid_path(id):
    return f'./tablescene/DiningRoom_{id}/frames/Objects/camera_0/Objects_0_0_0001_0.json'

def get_mask_path(id):
    return f'./tablescene/DiningRoom_{id}/frames/ObjectSegmentation/camera_0/ObjectSegmentation_0_0_0001_0.npy'

def get_scene_path_(path):
    return f'{path}/coarse/scene.blend'

def get_image_path_(path):
    return f'{path}/frames/Image/camera_0/Image_0_0_0001_0.png'

def get_depth_path_(path):
    return f'{path}/frames/Depth/camera_0/Depth_0_0_0001_0.npy'

def get_objid_path_(path):
    return f'{path}/frames/Objects/camera_0/Objects_0_0_0001_0.json'

def get_mask_path_(path):
    return f'{path}/frames/ObjectSegmentation/camera_0/ObjectSegmentation_0_0_0001_0.npy'

def random_scene():
    while True:
        scene_id = to_hex(randint(0, 4095))
        scene_path = get_scene_path(scene_id)
        try:
            with bpy.data.libraries.load(scene_path) as (data_from, data_to):
                object_names = data_from.objects
            found_table = False
            found_lamp = False
            for name in object_names:
                if "TableDiningFactory" in name:
                    found_table = True
                if "PointLampFactory" in name:
                    found_lamp = True
            if found_table and found_lamp:
                return scene_id
            print(f'Table or Lamp not found in scene {scene_id}')
        except Exception as e:
            print(f'Error loading scene {scene_id}: {e}')

def read_category_list():
    # open myobj/Xcate.json
    with open('qagen/Xcate.json', 'r') as f:
        data = json.load(f)
    return data

def read_obj_info():
    # open myobj/Xassets.json
    with open('qagen/Xassets.json', 'r') as f:
        data = json.load(f)
    return data

def read_color_info():
    # open myobj/Xcolor.json
    with open('qagen/Xcolor.json', 'r') as f:
        data = json.load(f)
    return data

def getAngle(x, y):
    return math.atan2(y, x)

def render(scene_id):
    render_command = [
        "python", "-m", "infinigen_examples.generate_nature",
        "--task=render",
        f"--input_folder={output_folder_prefix}/{scene_type}_{scene_id}/coarse",
        f"--output_folder={output_folder_prefix}/{scene_type}_{scene_id}/frames",
    ]
    
    flat_render_command = [
        "python", "-m", "infinigen_examples.generate_nature",
        "--task=render",
        f"--input_folder={output_folder_prefix}/{scene_type}_{scene_id}/coarse",
        f"--output_folder={output_folder_prefix}/{scene_type}_{scene_id}/frames",
        "-p", "render.render_image_func=@flat/render_image"
    ]
    
    subprocess.run(render_command)
    subprocess.run(flat_render_command)

def render_(scene_path):
    scene_path = os.path.expanduser(scene_path)
    render_command = [
        "python", "-m", "infinigen_examples.generate_nature",
        "--task=render",
        f"--input_folder={scene_path}/coarse",
        f"--output_folder={scene_path}/frames",
    ]
    
    flat_render_command = [
        "python", "-m", "infinigen_examples.generate_nature",
        "--task=render",
        f"--input_folder={scene_path}/coarse",
        f"--output_folder={scene_path}/frames",
        "-p", "render.render_image_func=@flat/render_image"
    ]
    
    subprocess.run(render_command)
    subprocess.run(flat_render_command)

def get_object_mask_rle(json_data, mask, object_name):
    """
    根据物体名称，从json_data中递归获取该物体以及所有子孙节点的object_index集合，
    并在mask中筛选出这些标号对应的区域，最终对区域进行RLE压缩后返回结果。

    参数：
    - json_data: dict，包含各物体信息的JSON数据，每个键为物体名字，值为属性字典。
      每个属性字典中包含 "object_index"（标号）以及可能的 "children"（子节点标号列表）等信息。
    - mask: numpy.ndarray，二维数组，每个像素点的值表示一个物体材质的标号。
    - object_name: str，起始物体的名字，从该节点开始搜索所有子孙节点。

    返回：
    - dict，包含两个键："counts" 表示 RLE 压缩后的计数列表，
      "size" 表示 mask 的形状（高度和宽度）。
    """
    
    # 1. 建立 object_index 到节点名的映射
    index_to_name = {}
    for name, node in json_data.items():
        if "object_index" in node:
            index_to_name[node["object_index"]] = name

    # 2. 递归搜集物体及其子孙节点的 object_index
    def collect_indices(name):
        node = json_data.get(name)
        if node is None:
            return set()
        
        # 将当前节点的 object_index 加入集合
        indices = {node.get("object_index")}
        
        # 遍历children（如果有的话），通过映射找到对应的节点名，递归收集标号
        for child_id in node.get("children", []):
            child_name = index_to_name.get(child_id)
            if child_name is not None:
                indices.update(collect_indices(child_name))
        return indices

    object_indices = collect_indices(object_name)

    # 3. 根据标号集合，从mask中筛选出对应的区域，生成二值mask
    #    np.isin返回的结果为布尔数组，True表示该像素属于物体区域
    binary_mask = np.isin(mask, list(object_indices))

    # 4. 对二值 mask 进行 RLE 压缩
    #    我们将mask打平成一维数组，然后计算连续相同值的运行长度
    flat_mask = binary_mask.flatten()
    rle_counts = []
    pref = None
    if flat_mask.size == 0:
        rle_counts = []
    else:
        # 初始化
        prev = flat_mask[0]
        count = 1
        pref = 1 if prev else 0
        # 从第二个元素开始遍历
        for pixel in flat_mask[1:]:
            if pixel == prev:
                count += 1
            else:
                rle_counts.append(count)
                count = 1
                prev = pixel
        rle_counts.append(count)

    return {"counts": rle_counts, "size": list(binary_mask.shape), "prefix": pref}
